import os
import shutil

Images_path = 'D:\\Tools\\yolo\\dlb\\PCB_DATASET\\images'  # 源图路径

train_labels = 'D:\\Tools\\yolo\\dlb\\PCB_DATASET\\labels\\train2017'  # train标签路径
val_labels = 'D:\\Tools\\yolo\\dlb\\PCB_DATASET\\labels\\val2017'  # val标签路径

train_images = 'D:\\Tools\\yolo\\dlb\\PCB_DATASET\\images\\train2017'  # 保存train图像路径
val_images = 'D:\\Tools\\yolo\\dlb\\PCB_DATASET\\images\\val2017'  # 保存val图像路径


# 判断文件夹是否存在，不存在即创建
if not os.path.exists(train_images):
    os.mkdir(train_images)
if not os.path.exists(val_images):
    os.mkdir(val_images)


# 按照标签名移动对应图像
for label_name in os.listdir(train_labels):
    img_name = label_name[:-3] + 'jpg'  # txt2jpg
    shutil.move(os.path.join(Images_path, img_name), os.path.join(train_images, img_name))


for label_name in os.listdir(val_labels):
    img_name = label_name[:-3] + 'jpg'  # txt2jpg
    shutil.move(os.path.join(Images_path, img_name), os.path.join(val_images, img_name))

